import torch
import torch.nn as nn

from torchvision.transforms import Compose, RandomCrop
from .rrc import RandomResizeCrop
from .mixup import LogMixupExp
from .byola import RandomCrop
from datasets.dataio import TorchScaler


class ATSTFrameTransform(nn.Module):
    def __init__(self):
        super().__init__()
        fw = RandomResizeCrop(virtual_crop_scale=(1.0, 1.0), freq_scale=(0.6, 1.5), time_scale=(1.0, 1.0))
        self.transform_1 = Compose([LogMixupExp(), fw])
        self.transform_2 = Compose([LogMixupExp(), fw])
        self.post_norm = TorchScaler("instance", "minmax", [1, 2])        

    def forward(self, x):
        if len(x.shape) == 3:
            view_1 = []
            view_2 = []
            for sample in x:
                sample = sample.unsqueeze(0)
                view_1.append(self.transform_1(sample))
                view_2.append(self.transform_2(sample))
            view_1 = torch.cat(view_1, dim=0)
            view_2 = torch.cat(view_2, dim=0)
            view_1 = self.post_norm(view_1)
            view_2 = self.post_norm(view_2)
        elif len(x.shape) == 2:
            x = x.unsqueeze(0)
            view_1 = self.transform_1(x)
            view_2 = self.transform_1(x)
            view_1 = self.post_norm(view_1)
            view_2 = self.post_norm(view_2)
            view_1 = view_1.squeeze(0)
            view_2 = view_2.squeeze(0)
        return view_1, view_2